'''
 # @ Author: Jeff Chen
 # @ Create Time: 2025-10-13 15:41:53
 # @ Description: 热成像处理核心类
 '''
import cv2
import numpy as np
import yaml
from sklearn.cluster import KMeans
from scipy import ndimage
from typing import Tuple, Dict, Any
import os

class ThermalImageProcessor:
    """
    热成像图像处理器 - 使用K-means聚类标注低温区域
    """
    
    def __init__(self, config_path: str = "config/default_config.yaml"):
        """
        初始化处理器
        
        Args:
            config_path: 配置文件路径
        """
        self.load_config(config_path)
        self.kmeans = None
        
    def load_config(self, config_path: str):
        """加载配置文件"""
        try:
            with open(config_path, 'r', encoding='utf-8') as file:
                self.config = yaml.safe_load(file)
        except FileNotFoundError:
            print(f"配置文件 {config_path} 未找到，使用默认配置")
            self.config = self.get_default_config()
    
    def get_default_config(self) -> Dict[str, Any]:
        """获取默认配置"""
        return {
            'kmeans': {
                'n_clusters': 3,
                'random_state': 42,
                'n_init': 10,
                'max_iter': 300
            },
            'preprocessing': {
                'blur_kernel_size': 5,
                'use_spatial_features': True,
                'spatial_weight': 0.1
            },
            'postprocessing': {
                'morphology_operation': True,
                'kernel_size': 3,
                'min_area': 100,
                'remove_small_objects': True
            }
        }
    
    def load_image(self, image_path: str) -> np.ndarray:
        """
        加载热成像图像
        
        Args:
            image_path: 图像路径
            
        Returns:
            加载的图像数组
        """
        if not os.path.exists(image_path):
            raise FileNotFoundError(f"图像文件 {image_path} 不存在")
        
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"无法加载图像 {image_path}")
            
        return image
    
    def preprocess_image(self, image: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        图像预处理
        
        Args:
            image: 输入图像
            
        Returns:
            gray_image: 灰度图像
            features: 特征矩阵
        """
        # 转换为灰度图
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        
        # 高斯模糊去噪
        blur_kernel = self.config['preprocessing']['blur_kernel_size']
        if blur_kernel > 0:
            gray = cv2.GaussianBlur(gray, (blur_kernel, blur_kernel), 0)
        
        # 准备特征矩阵
        pixels = gray.reshape(-1, 1).astype(np.float32)
        
        if self.config['preprocessing']['use_spatial_features']:
            # 添加空间位置特征
            height, width = gray.shape
            y_coords, x_coords = np.mgrid[0:height, 0:width]
            spatial_weight = self.config['preprocessing']['spatial_weight']
            coords = np.column_stack((x_coords.ravel(), y_coords.ravel())) * spatial_weight
            features = np.column_stack((pixels, coords))
        else:
            features = pixels
            
        return gray, features
    
    def apply_kmeans(self, features: np.ndarray) -> np.ndarray:
        """
        应用K-means聚类
        
        Args:
            features: 特征矩阵
            
        Returns:
            labels: 聚类标签
        """
        kmeans_config = self.config['kmeans']
        
        self.kmeans = KMeans(
            n_clusters=kmeans_config['n_clusters'],
            random_state=kmeans_config['random_state'],
            n_init=kmeans_config['n_init'],
            max_iter=kmeans_config['max_iter']
        )
        
        labels = self.kmeans.fit_predict(features)
        return labels
    
    def identify_low_temperature_cluster(self, labels: np.ndarray, 
                                       image_shape: Tuple[int, int]) -> np.ndarray:
        """
        识别低温聚类并生成掩码
        
        Args:
            labels: 聚类标签
            image_shape: 图像形状
            
        Returns:
            low_temp_mask: 低温区域掩码
        """
        if self.kmeans is None:
            raise ValueError("请先执行K-means聚类")
        
        # 获取聚类中心
        centers = self.kmeans.cluster_centers_
        if centers.shape[1] > 1:  # 如果有空间特征，只取温度特征
            temp_centers = centers[:, 0]
        else:
            temp_centers = centers.flatten()
        
        # 找到温度最低的聚类
        low_temp_cluster = np.argmin(temp_centers)
        
        # 生成低温区域掩码
        low_temp_mask = (labels == low_temp_cluster).reshape(image_shape)
        
        return low_temp_mask.astype(np.uint8)
    
    def postprocess_mask(self, mask: np.ndarray) -> np.ndarray:
        """
        后处理掩码
        
        Args:
            mask: 原始掩码
            
        Returns:
            processed_mask: 处理后的掩码
        """
        post_config = self.config['postprocessing']
        
        if post_config['morphology_operation']:
            # 形态学操作
            kernel_size = post_config['kernel_size']
            kernel = np.ones((kernel_size, kernel_size), np.uint8)
            
            # 开运算去除噪声
            mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
            # 闭运算填充空洞
            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
        
        if post_config['remove_small_objects']:
            # 去除小面积区域
            min_area = post_config['min_area']
            labeled_mask, num_features = ndimage.label(mask)
            
            for i in range(1, num_features + 1):
                component_mask = (labeled_mask == i)
                if np.sum(component_mask) < min_area:
                    mask[component_mask] = 0
        
        return mask
    
    def get_cluster_info(self) -> Dict[str, Any]:
        """
        获取聚类信息
        
        Returns:
            聚类统计信息字典
        """
        if self.kmeans is None:
            return {}
        
        centers = self.kmeans.cluster_centers_
        if centers.shape[1] > 1:
            temp_centers = centers[:, 0]
        else:
            temp_centers = centers.flatten()
        
        low_temp_cluster = np.argmin(temp_centers)
        
        return {
            'cluster_centers': temp_centers.tolist(),
            'low_temperature_cluster': int(low_temp_cluster),
            'low_temperature_value': float(temp_centers[low_temp_cluster])
        }
    
    def process_image(self, image_path: str) -> Dict[str, Any]:
        """
        完整处理流程
        
        Args:
            image_path: 输入图像路径
            
        Returns:
            处理结果字典
        """
        # 加载图像
        image = self.load_image(image_path)
        
        # 预处理
        gray, features = self.preprocess_image(image)
        
        # K-means聚类
        labels = self.apply_kmeans(features)
        
        # 识别低温区域
        low_temp_mask = self.identify_low_temperature_cluster(labels, gray.shape)
        
        # 后处理
        processed_mask = self.postprocess_mask(low_temp_mask)
        
        # 获取聚类信息
        cluster_info = self.get_cluster_info()
        
        results = {
            'original_image': image,
            'gray_image': gray,
            'low_temp_mask': processed_mask,
            'cluster_info': cluster_info,
            'features': features,
            'labels': labels
        }
        
        return results
